import torch
import torch.nn.functional as F

from lib.utils.one_hot import expand_as_one_hot

def structure_loss(pred, pred_bg, mask_fg, mask_bg, num_classes):
    if pred.shape != mask_fg.shape:
        mask_fg = expand_as_one_hot(mask_fg.long(), num_classes)
        mask_bg = expand_as_one_hot(mask_bg.long(), num_classes)

    # Weight (average pooling to get a smoothed version, calculate the absolute difference with the original mask, multiply by 5 to amplify the penalty for boundary differences, and finally add 1 to ensure the minimum weight is 1)
    weit = 1 + 5*torch.abs(F.avg_pool2d(mask_fg, kernel_size=31, stride=1, padding=15) - mask_fg) # Add 1 to ensure the minimum weight is 1

    # Weighted binary cross-entropy loss (foreground)
    wbce = F.binary_cross_entropy_with_logits(pred, mask_fg, reduction='none')
    wbce = (weit*wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3)) # 归一化
    
    # Weighted binary cross-entropy loss (background)
    wbce2 = F.binary_cross_entropy_with_logits(pred_bg, mask_bg, reduction='none')
    wbce2 = (weit*wbce2).sum(dim=(2, 3)) / weit.sum(dim=(2, 3)) # 归一化

    pred = torch.sigmoid(pred)
    
    # Weighted intersection over union loss
    inter = ((pred * mask_fg)*weit).sum(dim=(2, 3))
    union = ((pred + mask_fg)*weit).sum(dim=(2, 3))
    wiou = 1 - (inter + 1)/(union - inter+1)
    
    return (wbce + wiou + 0.8*wbce2).mean()